


from jittor.dataset import Dataset
import os
import cv2
import numpy as np

# IMG_SIZE=(224,224)
# IMG_SIZE=(299,299)
IMG_SIZE=(448,448)

import cv2


def resize_keep_aspectratio(image_src, dst_size):
    src_h, src_w = image_src.shape[:2]
    # print(src_h, src_w)
    dst_h, dst_w = dst_size

    # 判断应该按哪个边做等比缩放
    h = dst_w * (float(src_h) / src_w)  # 按照ｗ做等比缩放
    w = dst_h * (float(src_w) / src_h)  # 按照h做等比缩放

    h = int(h)
    w = int(w)

    if h <= dst_h:
        image_dst = cv2.resize(image_src, (dst_w, int(h)))
    else:
        image_dst = cv2.resize(image_src, (int(w), dst_h))

    h_, w_ = image_dst.shape[:2]
    # print(h_, w_)

    top = int((dst_h - h_) / 2);
    down = int((dst_h - h_ + 1) / 2);
    left = int((dst_w - w_) / 2);
    right = int((dst_w - w_ + 1) / 2);

    value = [0, 0, 0]
    borderType = cv2.BORDER_CONSTANT
    # print(top, down, left, right)
    image_dst = cv2.copyMakeBorder(image_dst, top, down, left, right, borderType, None, value)

    return image_dst

class Dogs(Dataset):
    def __init__(self, data_root="/home/hfle/data_dog", train=True , batch_size=1, shuffle=False):
        # if you want to test resnet etc you should set input_channel = 3, because the net set 3 as the input dimensions
        super().__init__()
        self.data_root = data_root
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.is_train = train
        
        lst_paths=(os.path.join(data_root,"train.lst"),(os.path.join(data_root,"validation.lst")))
        img_path=os.path.join(data_root,'low-resolution')
        # label map
        label_names=os.listdir(img_path)
        self.label_map={}
        for n in label_names:
            num=n.split('-')[1][-3:]
            self.label_map[n]=int(num)
        
        max_num=500
        num_dict=dict()
        if self.is_train:
            # train_data
            ft=open(lst_paths[0])
            train_names=ft.read().replace('.//','').rstrip('\n').split('\n')
            # train_names=train_names.replace('//','/low-resolution/').rstrip('\n').split('\n')
            self.imgs=list()
            self.labels=list()
            for i,n in enumerate(train_names):
                cls_n=n.split('/')[0]
                if cls_n not in num_dict.keys():
                    num_dict[cls_n]=1
                elif num_dict[cls_n]>=max_num:
                    continue
                else:
                    num_dict[cls_n]+=1
                    
                file_path=os.path.join(img_path,n)
                img=cv2.imread(file_path)
                img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
                # img=cv2.resize(img,(IMG_SIZE[0],IMG_SIZE[1]))
                img=resize_keep_aspectratio(img,IMG_SIZE)
                img=img.transpose((2,0,1))
                # img=img.astype(np.float32)
                self.imgs.append(img)
                self.labels.append(self.label_map[n.split('/')[0]])
                # print(i)
        else:
            # validation_data
            fv=open(lst_paths[1])
            val_names=fv.read().replace('.//','').rstrip('\n').split('\n')
            self.imgs=list()
            self.labels=list()
            for i,n in enumerate(val_names):
                # if n not in num_dict.keys():
                #     num_dict[n]=1
                # elif num_dict[n]>=max_num/10:
                #     continue
                # else:
                #     num_dict[n]+=1

                file_path=os.path.join(img_path,n)
                img=cv2.imread(file_path)
                img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
                # img=cv2.resize(img,(IMG_SIZE[0],IMG_SIZE[1]))
                img=resize_keep_aspectratio(img,IMG_SIZE)
                img=img.transpose((2,0,1))
                # img=img.astype(np.float32)
                self.imgs.append(img)
                self.labels.append(self.label_map[n.split('/')[0]])

        self.total_len=len(self.imgs)
        # this function must be called
        self.set_attrs(batch_size = self.batch_size, total_len=self.total_len, shuffle= self.shuffle)


    def __getitem__(self, index):
        return self.imgs[index].astype(np.float32)/255.0, self.labels[index]